//
//  OpenHaystack – Tracking personal Bluetooth devices via Apple's Find My network
//
//  Copyright © 2021 Secure Mobile Networking Lab (SEEMOO)
//  Copyright © 2021 The Open Wireless Link Project
//
//  SPDX-License-Identifier: AGPL-3.0-only
//

import CryptoKit
import Foundation
import OSLog

struct FindMyKeyExtractor {
  // swiftlint:disable identifier_name

  /// This function reads the private keys of the Offline Finding Location system. They will
  /// - Throws: Error when accessing files fails
  /// - Returns: Devices and their respective keys
  static func readPrivateKeys() throws -> [FindMyDevice] {
    var devices = [FindMyDevice]()
    os_log(.debug, "Looking for keys")

    do {

      // The key files have moved with macOS 10.15.4
      let macOS10_15_3Devices = try self.readFromOldLocation()
      devices.append(contentsOf: macOS10_15_3Devices)
    } catch {
      os_log(.error, "Did not find keys for 10.15.3\n%@", String(describing: error))
    }

    do {
      // Tries to discover the new location of the keys
      let macOS10_15_4Devices = try self.findKeyFilesInNewLocation()
      devices.append(contentsOf: macOS10_15_4Devices)
    } catch {
      os_log(.error, "Did not find keys for 10.15.4\n%@", String(describing: error))
    }

    return devices
  }

  // MARK: - macOS 10.15.0 - 10.15.3

  /// Reads the find my keys from the location used until macOS 10.15.3
  /// - Throws: An error if the location is no longer available (e.g. in macOS 10.15.4)
  /// - Returns: An array of find my devices including their keys
  static func readFromOldLocation() throws -> [FindMyDevice] {
    // Access the find my directory where the private advertisement keys are stored unencrypted
    let directoryPath = "com.apple.icloud.searchpartyd/PrivateAdvertisementKeys/"

    let fm = FileManager.default
    let privateKeysPath = fm.urls(for: .libraryDirectory, in: .userDomainMask)
      .first?.appendingPathComponent(directoryPath)
    let folders = try fm.contentsOfDirectory(
      at: privateKeysPath!,
      includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
    guard folders.isEmpty == false else { throw FindMyError.noFoldersFound }

    print("Found \(folders.count) folders")
    var devices = [FindMyDevice]()

    for folderURL in folders {
      let keyFiles = try fm.contentsOfDirectory(
        at: folderURL,
        includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
      // Check if keys are available
      print("Found \(keyFiles.count) in folder \(folderURL.lastPathComponent)")
      guard keyFiles.isEmpty == false else { continue }
      var device = FindMyDevice(deviceId: folderURL.lastPathComponent)

      for url in keyFiles {
        do {
          if url.pathExtension == "keys" {
            let keyPlist = try Data(contentsOf: url)
            let keyInfo = try self.parseKeyFile(keyFile: keyPlist)
            device.keys.append(keyInfo)
          }
        } catch {
          print("Could not load key file ", error)
        }

      }

      devices.append(device)
    }

    return devices
  }

  /// Parses the key plist file used until macOS 10.15.3
  /// - Parameter keyFile: Propery list data
  /// - Returns: Find My private Key
  static func parseKeyFile(keyFile: Data) throws -> FindMyKey {
    guard
      let keyDict = try PropertyListSerialization.propertyList(
        from: keyFile,
        options: .init(), format: nil) as? [String: Any],
      let advertisedKey = keyDict["A"] as? Data,
      let privateKey = keyDict["PR"] as? Data,
      let timeValues = keyDict["D"] as? [Double],
      let pu = keyDict["PU"] as? Data
    else {
      throw FindMyError.parsingFailed
    }

    let hashedKeyDigest = SHA256.hash(data: advertisedKey)
    let hashedKey = Data(hashedKeyDigest)
    let time = Date(timeIntervalSinceReferenceDate: timeValues[0])
    let duration = timeValues[1]

    return FindMyKey(
      advertisedKey: advertisedKey,
      hashedKey: hashedKey,
      privateKey: privateKey,
      startTime: time,
      duration: duration,
      pu: pu,
      yCoordinate: nil,
      fullKey: nil)
  }

  // MARK: - macOS 10.15.4 - 10.15.6 (+ Big Sur 11.0 Betas)

  /// Find the randomized key folder which is used since macOS 10.15.4
  /// - Returns: Returns an array of urls that contain keys. Multiple folders are found if the mac has multiple users
  static func findRamdomKeyFolder() -> [URL] {
    os_log(.debug, "Searching for cached keys folder")
    var folderURLs = [URL]()
    let foldersPath = "/private/var/folders/"
    let fm = FileManager.default

    func recursiveSearch(from url: URL, urlArray: inout [URL]) {
      do {
        let randomSubfolders = try fm.contentsOfDirectory(
          at: url,
          includingPropertiesForKeys: nil,
          options: .includesDirectoriesPostOrder)

        for folder in randomSubfolders {
          if folder.lastPathComponent == "com.apple.icloud.searchpartyd" {
            urlArray.append(folder.appendingPathComponent("Keys"))
            os_log(.debug, "Found folder at: %@", folder.path)
            break
          } else {
            recursiveSearch(from: folder, urlArray: &urlArray)
          }
        }

      } catch {

      }

    }

    recursiveSearch(from: URL(fileURLWithPath: foldersPath), urlArray: &folderURLs)

    return folderURLs

  }

  /// Find the key files in macOS 10.15.4 and newer (not working with fixed version 10.15.6)
  /// - Throws: An error if the key folder cannot be fould
  /// - Returns: An array of devices including their keys
  static func findKeyFilesInNewLocation() throws -> [FindMyDevice] {
    let keysFolders = self.findRamdomKeyFolder()
    guard keysFolders.isEmpty == false else {
      throw NSError(domain: "error", code: NSNotFound, userInfo: nil)
    }

    var devices = [FindMyDevice]()
    for folder in keysFolders {
      if let deviceKeys = try? self.loadNewKeyFilesIn(directory: folder) {
        devices.append(contentsOf: deviceKeys)
      }
    }

    return devices
  }

  /// Load the keys fils in the passed directory
  /// - Parameter directory: Pass a directory url to a location with key files
  /// - Throws: An error if the keys could not be found
  /// - Returns: An array of devices including their keys
  static func loadNewKeyFilesIn(directory: URL) throws -> [FindMyDevice] {
    os_log(.debug, "Loading key files from %@", directory.path)
    let fm = FileManager.default
    let subDirectories = try fm.contentsOfDirectory(
      at: directory,
      includingPropertiesForKeys: nil, options: .skipsHiddenFiles)

    var devices = [FindMyDevice]()

    for deviceDirectory in subDirectories {
      do {
        var keyFiles = [Data]()
        let keyDirectory = deviceDirectory.appendingPathComponent("Primary")
        let keyFileURLs = try fm.contentsOfDirectory(
          at: keyDirectory,
          includingPropertiesForKeys: nil,
          options: .skipsHiddenFiles)
        for keyfileURL in keyFileURLs {
          // Read the key files
          let keyFile = try Data(contentsOf: keyfileURL)
          if keyFile.isEmpty == false {
            keyFiles.append(keyFile)
          }
        }

        // Decode keys for file
        let decoder = FindMyKeyDecoder()
        var decodedKeys = [FindMyKey]()
        for file in keyFiles {
          do {
            let fmKeys = try decoder.parse(keyFile: file)
            decodedKeys.append(contentsOf: fmKeys)
          } catch {
            os_log(.error, "Decoding keys failed %@", error.localizedDescription)
          }
        }

        let device = FindMyDevice(deviceId: deviceDirectory.lastPathComponent, keys: decodedKeys)
        devices.append(device)
      } catch {
        os_log(.error, "Key directory not found %@", error.localizedDescription)
      }
    }

    return devices
  }

}
